"""
    ControlVectors

Module implementing control vector operations for GSI variational data assimilation.
This module is ported from the GSI Fortran `control_vectors.f90` module and provides
efficient implementations of vector operations, memory management, and ensemble handling.

The control vector represents the analysis increment in the control space, which is
related to the physical space through the background error covariance matrix B:
```
x = xₑ + B^(1/2) * δx
```

Where:
- `x`: Analysis in physical space
- `xₑ`: Background (first guess)  
- `δx`: Control vector increment
- `B^(1/2)`: Square root of background error covariance
"""
module ControlVectors

using LinearAlgebra
using Random
using ..GSICoreAnalysis: AbstractAnalysisConfig, AbstractControlVector, AnalysisVariables, default_analysis_variables

export ControlVector
export dot_product, axpy!, assign!, random_cv!, write_cv, read_cv
export norm_cv, maxval_cv, allocate_cv, deallocate_cv

"""
    ControlVector{T<:AbstractFloat}

Control vector for GSI variational data assimilation.

# Fields
- `values::Vector{T}`: Main control vector values
- `bias_predictors::Vector{T}`: Bias correction predictors  
- `ensemble_control::Vector{T}`: Ensemble control variables (for hybrid methods)
- `time_derivatives::Vector{T}`: Time derivative terms (for 4D-Var)
- `config::AbstractAnalysisConfig`: Analysis configuration
- `variables::AnalysisVariables`: Variable metadata
- `is_allocated::Bool`: Memory allocation status
"""
mutable struct ControlVector{T<:AbstractFloat} <: AbstractControlVector{T}
    values::Vector{T}
    bias_predictors::Vector{T}
    ensemble_control::Vector{T}
    time_derivatives::Vector{T}
    config::AbstractAnalysisConfig
    variables::AnalysisVariables
    is_allocated::Bool
    
    function ControlVector{T}(config::AbstractAnalysisConfig, 
                              variables::AnalysisVariables = default_analysis_variables()) where T
        cv = new{T}(
            T[],          # values
            T[],          # bias_predictors  
            T[],          # ensemble_control
            T[],          # time_derivatives
            config,       # config
            variables,    # variables
            false         # is_allocated
        )
        allocate_cv(cv)
        return cv
    end
end

# Convenience constructor
ControlVector(config::AbstractAnalysisConfig, args...; kwargs...) = 
    ControlVector{config.precision}(config, args...; kwargs...)

"""
    allocate_cv(cv::ControlVector{T})

Allocate memory for control vector components based on configuration.
"""
function allocate_cv(cv::ControlVector{T}) where T
    if cv.is_allocated
        return cv
    end

    # Calculate dimensions
    grid_size = cv.config.grid_size
    nx, ny, nz = grid_size

    # For compatibility with cost function, use same size as state vector
    # State size = 4*(nx*ny*nz) + (nx*ny) for (T, u, v, q) + ps
    n_3d_fields = 4
    n_2d_fields = 1
    main_size = n_3d_fields * (nx * ny * nz) + n_2d_fields * (nx * ny)

    cv.values = zeros(T, main_size)

    # Bias predictors (simplified - skip for now)
    cv.bias_predictors = T[]

    # Ensemble control (only if hybrid method enabled)
    if cv.config.use_hybrid && cv.config.ensemble_size > 0
        cv.ensemble_control = zeros(T, cv.config.ensemble_size)
    else
        cv.ensemble_control = T[]
    end

    # Time derivatives (skip for now)
    cv.time_derivatives = T[]

    cv.is_allocated = true
    return cv
end

"""
    deallocate_cv(cv::ControlVector)

Deallocate memory for control vector components.
"""
function deallocate_cv(cv::ControlVector)
    if !cv.is_allocated
        return cv
    end
    
    resize!(cv.values, 0)
    resize!(cv.bias_predictors, 0)
    resize!(cv.ensemble_control, 0)
    resize!(cv.time_derivatives, 0)
    
    cv.is_allocated = false
    return cv
end

"""
    dot_product(cv1::ControlVector{T}, cv2::ControlVector{T}) where T

Compute dot product of two control vectors.
Equivalent to GSI's `qdot_prod_cv` function.
"""
function dot_product(cv1::ControlVector{T}, cv2::ControlVector{T}) where T
    @assert cv1.is_allocated && cv2.is_allocated "Control vectors must be allocated"

    result = zero(T)

    # Main values contribution
    result += dot(cv1.values, cv2.values)

    # Bias predictors contribution (skip if empty)
    if !isempty(cv1.bias_predictors) && !isempty(cv2.bias_predictors)
        result += dot(cv1.bias_predictors, cv2.bias_predictors)
    end

    # Ensemble control contribution
    if !isempty(cv1.ensemble_control) && !isempty(cv2.ensemble_control)
        result += dot(cv1.ensemble_control, cv2.ensemble_control)
    end

    # Time derivatives contribution (skip if empty)
    if !isempty(cv1.time_derivatives) && !isempty(cv2.time_derivatives)
        result += dot(cv1.time_derivatives, cv2.time_derivatives)
    end

    return result
end

"""
    axpy!(α::T, x::ControlVector{T}, y::ControlVector{T}) where T

Perform y = y + α*x operation on control vectors.
Equivalent to GSI's `axpy` subroutine.
"""
function axpy!(α::T, x::ControlVector{T}, y::ControlVector{T}) where T
    @assert x.is_allocated && y.is_allocated "Control vectors must be allocated"

    # Main values: y.values += α * x.values
    LinearAlgebra.axpy!(α, x.values, y.values)

    # Bias predictors (skip if empty)
    if !isempty(x.bias_predictors) && !isempty(y.bias_predictors)
        LinearAlgebra.axpy!(α, x.bias_predictors, y.bias_predictors)
    end

    # Ensemble control
    if !isempty(x.ensemble_control) && !isempty(y.ensemble_control)
        LinearAlgebra.axpy!(α, x.ensemble_control, y.ensemble_control)
    end

    # Time derivatives (skip if empty)
    if !isempty(x.time_derivatives) && !isempty(y.time_derivatives)
        LinearAlgebra.axpy!(α, x.time_derivatives, y.time_derivatives)
    end

    return y
end

"""
    assign!(dest::ControlVector{T}, src::ControlVector{T}) where T

Copy control vector values: dest = src.
Equivalent to GSI's `assign_cv2cv` subroutine.
"""
function assign!(dest::ControlVector{T}, src::ControlVector{T}) where T
    @assert dest.is_allocated && src.is_allocated "Control vectors must be allocated"

    copyto!(dest.values, src.values)

    # Bias predictors (skip if empty)
    if !isempty(src.bias_predictors) && !isempty(dest.bias_predictors)
        copyto!(dest.bias_predictors, src.bias_predictors)
    end

    if !isempty(src.ensemble_control) && !isempty(dest.ensemble_control)
        copyto!(dest.ensemble_control, src.ensemble_control)
    end

    # Time derivatives (skip if empty)
    if !isempty(src.time_derivatives) && !isempty(dest.time_derivatives)
        copyto!(dest.time_derivatives, src.time_derivatives)
    end

    return dest
end

"""
    random_cv!(cv::ControlVector{T}) where T

Fill control vector with random values.
Equivalent to GSI's `random_cv` subroutine.
"""
function random_cv!(cv::ControlVector{T}) where T
    @assert cv.is_allocated "Control vector must be allocated"

    # Generate Gaussian random values
    randn!(cv.values)

    if !isempty(cv.bias_predictors)
        randn!(cv.bias_predictors)
    end

    if !isempty(cv.ensemble_control)
        randn!(cv.ensemble_control)
    end

    if !isempty(cv.time_derivatives)
        randn!(cv.time_derivatives)
    end

    return cv
end

"""
    norm_cv(cv::ControlVector{T}) where T

Compute L2 norm of control vector.
"""
function norm_cv(cv::ControlVector{T}) where T
    return sqrt(dot_product(cv, cv))
end

"""
    maxval_cv(cv::ControlVector{T}) where T

Compute maximum absolute value in control vector.
Equivalent to GSI's `maxval_cv` function.
"""
function maxval_cv(cv::ControlVector{T}) where T
    @assert cv.is_allocated "Control vector must be allocated"
    
    max_val = maximum(abs, cv.values)
    max_val = max(max_val, maximum(abs, cv.bias_predictors))
    
    if !isempty(cv.ensemble_control)
        max_val = max(max_val, maximum(abs, cv.ensemble_control))
    end
    
    max_val = max(max_val, maximum(abs, cv.time_derivatives))
    
    return max_val
end

"""
    write_cv(cv::ControlVector{T}, filename::String) where T

Write control vector to file.
Equivalent to GSI's `write_cv` subroutine.
"""
function write_cv(cv::ControlVector{T}, filename::String) where T
    @assert cv.is_allocated "Control vector must be allocated"
    
    open(filename, "w") do io
        # Write header information
        println(io, "# GSI Control Vector")
        println(io, "# Precision: $(T)")
        println(io, "# Main values size: $(length(cv.values))")
        println(io, "# Bias predictors size: $(length(cv.bias_predictors))")
        println(io, "# Ensemble control size: $(length(cv.ensemble_control))")
        println(io, "# Time derivatives size: $(length(cv.time_derivatives))")
        println(io)
        
        # Write data sections
        println(io, "# Main values")
        for val in cv.values
            println(io, val)
        end
        
        println(io, "# Bias predictors")
        for val in cv.bias_predictors
            println(io, val)
        end
        
        if !isempty(cv.ensemble_control)
            println(io, "# Ensemble control")
            for val in cv.ensemble_control
                println(io, val)
            end
        end
        
        println(io, "# Time derivatives")
        for val in cv.time_derivatives
            println(io, val)
        end
    end
    
    return nothing
end

"""
    read_cv(cv::ControlVector{T}, filename::String) where T

Read control vector from file.
Equivalent to GSI's `read_cv` subroutine.
"""
function read_cv(cv::ControlVector{T}, filename::String) where T
    @assert cv.is_allocated "Control vector must be allocated"
    
    # Simple implementation - would need more robust parsing
    lines = readlines(filename)
    data_lines = filter(line -> !startswith(line, "#") && !isempty(strip(line)), lines)
    
    # Parse numeric data
    values = [parse(T, line) for line in data_lines]
    
    # Distribute values to appropriate arrays (simplified)
    n_main = length(cv.values)
    n_bias = length(cv.bias_predictors)
    n_ens = length(cv.ensemble_control)
    n_time = length(cv.time_derivatives)
    
    idx = 1
    if n_main > 0
        copyto!(cv.values, values[idx:idx+n_main-1])
        idx += n_main
    end
    
    if n_bias > 0
        copyto!(cv.bias_predictors, values[idx:idx+n_bias-1])
        idx += n_bias
    end
    
    if n_ens > 0
        copyto!(cv.ensemble_control, values[idx:idx+n_ens-1])
        idx += n_ens
    end
    
    if n_time > 0
        copyto!(cv.time_derivatives, values[idx:idx+n_time-1])
    end
    
    return cv
end

end # module ControlVectors